﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Mono.Cecil;
using Mono.Cecil.Cil;
using NETDeob.Engine.Utils;
using NETDeob.Misc;
using NETDeob.Misc.Structs__Enums___Interfaces.Deobfuscation;
using NETDeob.Misc.Structs__Enums___Interfaces.Tasks.Base;

namespace NETDeob.Deobfuscators.Obfusasm.Tasks
{
    class StringDecryptor : DeobfuscationTask
    {
        private class ObfusasmEntry : DecryptionEntry
        {
            public int Key;
            public int Key1;
            public int Key2;

            public List<Tuple<Instruction, MethodDefinition>> InsertTargets = new List<Tuple<Instruction,MethodDefinition>>();
            public MethodDefinition TargetProxy;

            public override string ToString()
            {
                return string.Format(@"[Decrypt]({4}) Key({0}, {1}, {2}) -> ""{3}""", Key, Key1, Key2, PlainText, InsertTargets.Count);
            }
        }

        public StringDecryptor(AssemblyDefinition asmDef)
            : base(asmDef)
        {
            RoutineDescription = "String decryption";
        }

        private static byte[] _staticKey;

        [DeobfuscationPhase(1, "Locate bad type")]
        public static bool Phase1()
        {
            MethodDefinition cctor = null;
            var targetType = AsmDef.MainModule.Types.FirstOrDefault(t => IsBadType(t, out cctor));

            if(targetType == null || cctor == null)
            {
                PhaseError = new PhaseError
                                 {
                                     Level = PhaseError.ErrorLevel.Minor,
                                     Message = "No string encryption?"
                                 };
                EmergencyCancel = true;
                return true;
            }

            MarkMember(targetType);
            Logger.VSLog("Located bad type at " + targetType.Name);

            _staticKey = (cctor.Body.Instructions.FirstOrDefault(op => op.OpCode == OpCodes.Ldtoken).Operand as FieldReference).Resolve().InitialValue;

            if (_staticKey == null || _staticKey.Length == 0){
                ThrowPhaseError("Could not find static key!", 0, true);
                return true;
            }

            PhaseParam = targetType; 

            return true;
        }

        /// <summary>
        /// PhaseParam = targetType
        /// </summary>
        [DeobfuscationPhase(2, "Construct decryption entries")]
        public static bool Phase2()
        {
            var targetType = PhaseParam;
            var decEntries = new List<ObfusasmEntry>();
            var usedMethods = new List<MethodDefinition>();

            foreach(var mDef in AsmDef.FindMethods(m => true))
            {
                if (!mDef.HasBody)
                    continue;

                for (var i = 0; i < mDef.Body.Instructions.Count; i++)
                {
                    var instr = mDef.Body.Instructions[i];
                    MethodDefinition tmpTarget = null;

                    if (instr.OpCode == OpCodes.Call &&
                        (tmpTarget = (instr.Operand as MethodReference).Resolve()).DeclaringType == targetType)
                    {
                        if (!IsStringProxy(tmpTarget))
                            continue;

                        var decEntry = new ObfusasmEntry
                                           {
                                               Key =
                                                   tmpTarget.Body.Instructions.FindInstruction(
                                                       iinstr => iinstr.IsLdcI4(), 2).GetLdcI4(),
                                               Key1 =
                                                   tmpTarget.Body.Instructions.FindInstruction(
                                                       iinstr => iinstr.IsLdcI4(), 3).GetLdcI4(),
                                               Key2 =
                                                   tmpTarget.Body.Instructions.FindInstruction(
                                                       iinstr => iinstr.IsLdcI4(), 4).GetLdcI4(),
                                               TargetProxy = tmpTarget
                                           };

                        decEntry.InsertTargets.Add(new Tuple<Instruction,MethodDefinition>(instr.Next, mDef));
                        MarkMember(instr, mDef);

                        if (!usedMethods.Contains(tmpTarget)){
                            decEntries.Add(decEntry);
                            usedMethods.Add(tmpTarget);
                        }
                        else
                        {
                            decEntries.First(de => de.TargetProxy == tmpTarget).InsertTargets.Add(new Tuple<Instruction,MethodDefinition>(instr, mDef));
                            MarkMember(instr, mDef);
                        }
                    }
                }
            }

            PhaseParam = decEntries;

            return true;
        }

        /// <summary>
        /// PhaseParam = decEntries
        /// </summary>
        [DeobfuscationPhase(3, "Decrypt strings")]
        public static bool Phase3()
        {
            var decEntries = PhaseParam as List<ObfusasmEntry>;

           for (int i = 0; i < decEntries.Count; i++)
           {
               var entry = decEntries[i];
               DecryptEntry(ref entry);

               Logger.VLog(entry.ToString());
           }

            Logger.VSLog(string.Format("{0} strings decrypted...", decEntries.Count));
            return true;
        }

        /// <summary>
        /// PhaseParam = decEntries
        /// </summary>
        [DeobfuscationPhase(4, "Replace strings")]
        public static bool Phase4()
        {
            var decEntries = PhaseParam as List<ObfusasmEntry>;

            foreach (var entry in decEntries)
            {
                foreach (var target in entry.InsertTargets)
                {
                    var ilProc = target.Item2.Body.GetILProcessor();
                    ilProc.InsertBefore(target.Item1,
                                       ilProc.Create(OpCodes.Ldstr, entry.PlainText));
                }
            }

            Logger.VSLog(string.Format("{0} strings replaced in assembly with plaintext...", decEntries.Count));

            return true;
        }
   
        private static bool IsBadType(TypeDefinition typeDef, out MethodDefinition cctor)
        {
            if ((cctor = typeDef.GetStaticConstructor()) == null)
                return false;

            if (!cctor.HasBody)
                return false;

            if (cctor.Body.HasVariables || cctor.Body.Instructions.FirstOrDefault(i => i.OpCode == OpCodes.Ldtoken) == null)
                return false;

            return true;
        }
        private static bool IsStringProxy(MethodDefinition mDef)
        {
            if (mDef.Body.HasVariables || mDef.HasParameters)
                return false;

            if (mDef.ReturnType.Name != "String")
                return false;

            return true;
        }

        #region Reversed methods

        private static void DecryptEntry(ref ObfusasmEntry entry)
        {
            for (int i = entry.Key; i < (entry.Key + entry.Key1); i++)
                _staticKey[i] ^= (byte)entry.Key2;

            entry.PlainText = Encoding.UTF8.GetString(_staticKey, entry.Key, entry.Key1);
        }

        #endregion
    }
}
